import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader, random_split
from transformers import BertModel
from tqdm import tqdm
import numpy as np
from torch.utils.data import Dataset


# DisasterDataset 类定义
class DisasterDataset(Dataset):
    def __init__(self, file_path):
        data = np.load(file_path, allow_pickle=True).item()
        self.input_ids = data['input_ids']
        self.attention_mask = data['attention_mask']
        self.labels = data['labels']

    def __len__(self):
        return len(self.input_ids)

    def __getitem__(self, idx):
        return {
            'input_ids': self.input_ids[idx],
            'attention_mask': self.attention_mask[idx],
            'labels': self.labels[idx]
        }

# 定义 BertClassificationModel 类
class BertClassificationModel(nn.Module):
    def __init__(self, model_path, dropout=0.1):
        super(BertClassificationModel, self).__init__()
        self.bert = BertModel.from_pretrained('E:/bert-base')  # 使用本地模型路径
        self.dropout = nn.Dropout(dropout)
        self.fc = nn.Linear(self.bert.config.hidden_size, 2)

    def forward(self, input_ids, attention_mask):
        outputs = self.bert(input_ids=input_ids, attention_mask=attention_mask)
        pooled_output = outputs.pooler_output
        pooled_output = self.dropout(pooled_output)
        logits = self.fc(pooled_output)
        return logits

# 主训练过程
def main():
    batch_size = 256
    epoch_num = 2
    dropout = 0.3
    learning_rate = 2e-5

    # 加载数据集
    file_path = 'E:/nlp-getting-started/sentences_tokenizer.npy'
    dataset = DisasterDataset(file_path)

    # 将数据集划分为训练集和验证集
    train_size = 6144
    val_size = len(dataset) - train_size
    train_data, val_data = random_split(dataset, [train_size, val_size])

    # 定义数据加载器
    train_loader = DataLoader(train_data, batch_size=batch_size, shuffle=True)
    val_loader = DataLoader(val_data, batch_size=batch_size, shuffle=False)

    # 初始化模型和优化器
    model_path = 'E:/bert-base/'  # 指定本地模型路径
    model = BertClassificationModel(model_path, dropout=dropout)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model.to(device)
    if torch.cuda.device_count() > 1:
        model = nn.DataParallel(model)

    optimizer = optim.AdamW(model.parameters(), lr=learning_rate)
    criterion = nn.CrossEntropyLoss()

    # 训练过程
    print("模型数据已经加载完成,现在开始模型训练。")
    for epoch in range(epoch_num):
        model.train()
        total_loss = 0.0
        correct_predictions = 0
        total_samples = 0

        for batch in tqdm(train_loader, desc=f'Epoch {epoch + 1}/{epoch_num}'):
            input_ids = batch['input_ids'].to(device)
            attention_mask = batch['attention_mask'].to(device)
            labels = batch['labels'].to(device)

            optimizer.zero_grad()
            outputs = model(input_ids, attention_mask)
            _, predicted = torch.max(outputs, 1)
            correct_predictions += (predicted == labels).sum().item()
            total_samples += labels.size(0)

            loss = criterion(outputs, labels)
            total_loss += loss.item()

            loss.backward()
            optimizer.step()

        train_accuracy = correct_predictions / total_samples
        avg_loss = total_loss / len(train_loader)
        print(f'Epoch {epoch + 1}/{epoch_num}, Train Loss: {avg_loss:.4f}, Train Accuracy: {train_accuracy:.4f}')

    # 在验证集上评估模型
    model.eval()
    with torch.no_grad():
        correct_predictions = 0
        total_samples = 0

        for batch in val_loader:
            input_ids = batch['input_ids'].to(device)
            attention_mask = batch['attention_mask'].to(device)
            labels = batch['labels'].to(device)

            outputs = model(input_ids, attention_mask)
            _, predicted = torch.max(outputs, 1)
            correct_predictions += (predicted == labels).sum().item()
            total_samples += labels.size(0)

        val_accuracy = correct_predictions / total_samples
        print(f'Validation Accuracy: {val_accuracy:.4f}')

    # 保存模型
    torch.save(model.state_dict(), 'E:/nlp-getting-started/bert_model.pth')
    print("模型已保存。")

if __name__ == '__main__':
    main()